_base_ = 'retinanet_r18_fpn_1x_coco.py'
import os
# 重写训练轮次为150轮
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=15, val_interval=1)

# model配置 修改类别数为2；使用resnet18
model = dict(
    backbone=dict(
        depth=18,
        init_cfg=dict(type='Pretrained', checkpoint='torchvision://resnet18')),
    neck=dict(in_channels=[64, 128, 256, 512]),
    bbox_head=dict(
        num_classes=2  # 重写类别数为2
    )
)

# learning rate 在80和110epoch修改学习率
param_scheduler = [
    dict(
        type='LinearLR', start_factor=0.01, by_epoch=False, begin=0, end=500),
    dict(
        type='MultiStepLR',
        begin=0,  # 开始回合
        end=15,  # 结束回合
        milestones=[8, 11])
]

# 修改数据集相关配置（绝对路径）
data_root = '/root/mmdetection/dataset/SDB_9K_COCO/'
metainfo = {
    'classes': ('drone', 'bird'),
    'palette': [
        (253, 58, 52), (253, 159, 148),
    ]
}
# 三个loader
train_dataloader = dict(
    batch_size=4,
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='train/annotations/SDB_9K_train.json',
        data_prefix=dict(img='train/images')))
val_dataloader = dict(
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='val/annotations/SDB_9K_val.json',
        data_prefix=dict(img='val/images')))
test_dataloader = dict(
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='test/annotations/SDB_9K_test.json',
        data_prefix=dict(img='test/images')))

# 修改评价指标相关配置
val_evaluator = dict(ann_file=os.path.join(data_root, 'val/annotations/SDB_9K_val.json'))
test_evaluator = dict(ann_file=os.path.join(data_root, 'test/annotations/SDB_9K_test.json'))

load_from = 'https://download.openmmlab.com/mmdetection/v2.0/retinanet/retinanet_r18_fpn_1x_coco/retinanet_r18_fpn_1x_coco_20220407_171055-614fd399.pth'
